import requests
from typing import List
from langchain_core.embeddings import Embeddings
import tiktoken

class SiliconFlowEmbeddings(Embeddings):
    """硅基流动 Embedding 模型"""
    
    def __init__(self, api_key: str, model: str = "BAAI/bge-m3"):
        self.api_key = api_key
        self.model = model
        self.base_url = "https://api.siliconflow.cn/v1/embeddings"
        
        # 模型限制
        self.max_tokens = {
            "BAAI/bge-large-zh-v1.5": 512,
            "BAAI/bge-large-en-v1.5": 512,
            "netease-youdao/bce-embedding-base_v1": 512,
            "BAAI/bge-m3": 8192,
            "Pro/BAAI/bge-m3": 8192,
            "Qwen/Qwen3-Embedding-8B": 32768,
            "Qwen/Qwen3-Embedding-4B": 32768,
            "Qwen/Qwen3-Embedding-0.6B": 32768
        }
        
        self.max_batch_size = 32  # 硅基流动限制
        self.max_tokens_per_text = self.max_tokens.get(model, 512) 
        
        # 初始化tokenizer用于计算token数量
        try:
            self.tokenizer = tiktoken.get_encoding("cl100k_base")
        except:
            self.tokenizer = None
    
    def _truncate_text(self, text: str) -> str:
        """截断文本到最大token限制"""
        if not text or not text.strip():
            return "无内容"  # 避免空字符串
        
        if self.tokenizer:
            # 使用tokenizer精确计算
            tokens = self.tokenizer.encode(text)
            if len(tokens) > self.max_tokens_per_text:
                # 截断到最大长度
                truncated_tokens = tokens[:self.max_tokens_per_text - 100]  # 留一些余量
                text = self.tokenizer.decode(truncated_tokens)
                print(f"⚠️ 文本被截断到 {len(truncated_tokens)} tokens")
        else:
            # 简单字符截断（不太精确但安全）
            max_chars = self.max_tokens_per_text * 2  # 粗略估计
            if len(text) > max_chars:
                text = text[:max_chars]
                print(f"⚠️ 文本被截断到 {max_chars} 字符")
        
        return text.strip() or "无内容"
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """嵌入多个文档"""
        if not texts:
            return []
        
        embeddings = []
        
        # 预处理文本：截断和去空
        processed_texts = []
        for text in texts:
            processed_text = self._truncate_text(text)
            processed_texts.append(processed_text)
        
        # 按硅基流动的限制分批处理
        for i in range(0, len(processed_texts), self.max_batch_size):
            batch_texts = processed_texts[i:i + self.max_batch_size]
            print(f"🔄 处理批次 {i//self.max_batch_size + 1}: {len(batch_texts)} 个文本")
            
            batch_embeddings = self._get_embeddings(batch_texts)
            if batch_embeddings:
                embeddings.extend(batch_embeddings)
            else:
                # 如果批次失败，添加零向量占位
                print(f"⚠️ 批次 {i//self.max_batch_size + 1} 失败，使用零向量占位")
                zero_embeddings = [[0.0] * 1024 for _ in batch_texts]  # 假设1024维
                embeddings.extend(zero_embeddings)
        
        print(f"✅ 总共处理 {len(embeddings)} 个文档嵌入")
        return embeddings
    
    def embed_query(self, text: str) -> List[float]:
        """嵌入单个查询"""
        processed_text = self._truncate_text(text)
        embeddings = self._get_embeddings([processed_text])
        return embeddings[0] if embeddings else []
    
    def _get_embeddings(self, texts: List[str]) -> List[List[float]]:
        """调用硅基流动 API 获取 embeddings"""
        try:
            # 验证输入
            if not texts:
                print("⚠️ 输入文本列表为空")
                return []
            
            if len(texts) > self.max_batch_size:
                print(f"⚠️ 批次大小 {len(texts)} 超过限制 {self.max_batch_size}")
                texts = texts[:self.max_batch_size]
            
            # 确保没有空字符串
            valid_texts = []
            for text in texts:
                if text and text.strip():
                    valid_texts.append(text.strip())
                else:
                    valid_texts.append("无内容")
            
            print(f"📤 发送 {len(valid_texts)} 个文本到硅基流动API")
            
            payload = {
                "model": self.model,
                "input": valid_texts,
                "encoding_format": "float"
            }
            
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            
            response = requests.post(
                self.base_url,
                json=payload,
                headers=headers,
                timeout=120  # 增加超时时间
            )
            
            if response.status_code == 200:
                result = response.json()
                print(f"✅ API 调用成功，返回 {len(result['data'])} 个嵌入")
                
                # 按 index 排序确保顺序正确
                data = sorted(result["data"], key=lambda x: x["index"])
                embeddings = [item["embedding"] for item in data]
                
                # 验证嵌入维度
                if embeddings:
                    dim = len(embeddings[0])
                    print(f"📊 嵌入维度: {dim}")
                
                return embeddings
            else:
                print(f"❌ Embedding API 错误: {response.status_code}")
                print(f"错误详情: {response.text}")
                return []
                
        except requests.exceptions.ConnectionError as e:
            print(f"❌ 网络连接错误: {e}")
            print("🔍 请检查网络连接和API服务状态")
            return []
        except requests.exceptions.Timeout as e:
            print(f"❌ 请求超时: {e}")
            return []
        except Exception as e:
            print(f"❌ Embedding 调用失败: {e}")
            import traceback
            traceback.print_exc()
            return []
    
    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        """异步嵌入多个文档"""
        return self.embed_documents(texts)
    
    async def aembed_query(self, text: str) -> List[float]:
        """异步嵌入单个查询"""
        return self.embed_query(text)